/*
 *
 * Portions of code from PspLink:
 * Copyright (c) 2005 James F <tyranid@gmail.com>
 * Some portions edited by Mortimus - 2008
 */
#include <pspkernel.h>
#include <pspdebug.h>
#include <pspsysmem_kernel.h>
#include <psputilsforkernel.h>
#include <pspmoduleexport.h>
#include <psploadcore.h>
#include <pspthreadman_kernel.h>
#include <pspsdk.h>
#include <stdio.h>
#include <string.h>
#include "apihook.h"

#define DEBUG_APIHOOK 0

//void //sioPrintWord(u32 value);
//void //sioPrint(const char* str);
//void //sioPrintHex(char ch);

typedef struct HookRecord
{
 char modname[32];
 char libname[32];
 u32 nid;
 void *func;
 void *oldfunc;
 void *syscall;
} HookRecord;

#define MAX_HOOK_RECORDS 30

HookRecord hookRecords[30];
int        numHookRecords = 0;
SceUID     hookSema = 0;

struct SyscallHeader
{
	void *unk;
	unsigned int basenum;
	unsigned int topnum;
	unsigned int size;
};

void initApiHooks(void)
{
	hookSema = sceKernelCreateSema("pikey_hooks", 0, 1, 1, NULL);
	numHookRecords = 0;
}

void *find_syscall_addr(u32 addr)
{
	struct SyscallHeader *head;
	u32 *syscalls;
	void **ptr;
	int size;
	int i;
	char text[255];
	//char asem[14];
	//int reg = 8;
	//while(asem<29){
	//sprintf(asem, "cfc0 %%0, $%i\n", reg); 

	// 0 Crashes
	// 1 Null PTR
	// 2 Crashes
	// 3 Crashes
	// 4 Crashes
	// 5 
	// 6
	// 7
	// 8 works
	// 9 Crashes
	// 10 Null PTR
	// 11 Crashes
	// 12 Original
	// 13 Crashes
	// 14 Crashes
	// 15 Null PTR
	// 16 Crashes
	// 17 Crashes
	// 18
	// 19
	// 20

	asm(
			"cfc0 %0, $8\n"
			: "=r"(ptr)
	   );

	if(!ptr)
	{
		ShowOutput("Null ptr from cfc");
		////sioPrint("Null ptr from cfc\n");
		return NULL;
	}

	head = (struct SyscallHeader *) *ptr;
	syscalls = (u32*) (*ptr + 0x10); // *ptr + 16
	size = (head->size - 0x10) / sizeof(u32);
	//size *= 2;
	//size = (head->size - 0x10);

	for(i = 0; i < size; i++)
	{
		//sprintf(text, "HEX: %X DEC: %i == HEX: %X DEC: %i", syscalls[i], syscalls[i], addr, addr);
		//ShowOutput(text);
		//sceKernelDelayThread(2500000);
		if(syscalls[i] == addr)
		{
			return &syscalls[i];
		}
	}
	sprintf(text, "syscall not found - HEX: %X DEC: %i SIZE: %i", addr, addr, size);
	ShowOutput(text);
	////sioPrint("Failed to find syscall for addr\n");
	//asem++;
	//};
	return NULL;
}

static void *apiHookAddr(u32 *addr, void *func)
{
	int intc;

	if(!addr)
	{
		//ShowOutput("Address is NULL");
		return NULL;
	}

	intc = pspSdkDisableInterrupts();
	*addr = (u32) func;
	sceKernelDcacheWritebackInvalidateRange(addr, sizeof(addr));
	sceKernelIcacheInvalidateRange(addr, sizeof(addr));
	pspSdkEnableInterrupts(intc);

	return addr;
}

static struct SceLibraryEntryTable *_libsFindLibrary(SceUID uid, const char *library)
{
	struct SceLibraryEntryTable *entry;
	SceModule *pMod;
	void *entTab;
	int entLen;
	//char text[255];

	pMod = sceKernelFindModuleByUID(uid);
	//pMod = sceKernelFindModuleByName("sceNetAdhoc_Library");
	if(pMod != NULL)
	{
		int i = 0;

		entTab = pMod->ent_top;
		entLen = pMod->ent_size;
		while(i < entLen)
		{
			entry = (struct SceLibraryEntryTable *) (entTab + i);
			//sprintf(text, "Library: %s", entry->libname);
			//ShowOutput(text);
			//sceKernelDelayThread(2500000);

			if((entry->libname) && (strcmp(entry->libname, library) == 0))
			{
				return entry;
			}
			else if(!entry->libname && !library)
			{
				return entry;
			}

			i += (entry->len * 4);
		}
		ShowOutput("Can't find libname in module");
		////sioPrint("Can't find libname in module\n");
	}
	else
	{
		ShowOutput("Can't find module UID");
		////sioPrint("Can't find module UID\n");
	}

	return NULL;
}

void* libsFindExportAddrByNid(SceUID uid, const char *library, u32 nid)
{
	struct SceLibraryEntryTable *entry;
	u32 *addr = NULL;
	//char text[255];

	entry = _libsFindLibrary(uid, library);
	if(entry)
	{
		int count;
		int total;
		unsigned int *vars;

		total = entry->stubcount + entry->vstubcount;
		vars = entry->entrytable;

		if(entry->stubcount > 0)
		{
			for(count = 0; count < entry->stubcount; count++)
			{
				//sprintf(text, "HEX: %X DEC: %i :: HEX: %X DEC: %i", vars[count], vars[count], nid, nid);
				//ShowOutput(text);
				//sceKernelDelayThread(2500000);	
				if(vars[count] == nid)
				{
					return &vars[count+total];
				}
			}
		}
		else
		{
			ShowOutput("No stubs for entry");
			////sioPrint("No stubs for entry\n");
		}
	}
	else
	{
		ShowOutput("Cannot find library");
	};
	return addr;
}

u32 libsFindExportByNid(SceUID uid, const char *library, u32 nid)
{
	u32 *addr;

	addr = libsFindExportAddrByNid(uid, library, nid);
	if(!addr)
	{
		ShowOutput("Cannot libsFindExportAddrByNid");
		return 0;
	}

	return *addr;
}

u32 apiHookByNid2(SceUID mID, const char *modname, const char *library, u32 nid, void *func)
{
	sceKernelWaitSema(hookSema, 1, 0);

	if (numHookRecords == MAX_HOOK_RECORDS)
	{
		sceKernelSignalSema(hookSema, 1);
		return -1;
	}

	u32 lhookid = -1;

  /***************************************************************************/
	/* First, let's see if we already have a hook for this function.           */
	/***************************************************************************/
	int ii;
	for (ii = 0; ii < numHookRecords; ii++)
	{
		if ((hookRecords[ii].nid == nid) &&
				(strcmp(hookRecords[ii].libname, library) == 0) &&
				(strcmp(hookRecords[ii].modname, modname) == 0))
		{
			lhookid = ii;
		}
	}

	if (lhookid == -1)
	{
		// no existing hook, so this is simple.
		//SceModule *pMod = sceKernelFindModuleByName(modname);
			u32 addr = libsFindExportByNid(mID, library, nid);
			if(addr)
			{
				void * syscalladdr = find_syscall_addr(addr);
				if(!apiHookAddr(syscalladdr, func))
				{
					////sioPrint("Failed to apiHookAddr\n");
					ShowOutput("Failed to apiHookAddr2");
				}
				else
				{
					// success hooking.  Record the details.
					lhookid = numHookRecords++;
					hookRecords[lhookid].oldfunc = (void*)addr;
					strcpy(hookRecords[lhookid].libname, library);
					strcpy(hookRecords[lhookid].modname, modname);
					hookRecords[lhookid].nid     = nid;
					hookRecords[lhookid].func    = func;
					hookRecords[lhookid].syscall = syscalladdr;
				}
			}
			else
			{
				ShowOutput("Failed to find exportbynid");
				////sioPrint("Failed to find exportbynid\n");
			}
	}
	else
	{
		// We have an existing hook for this function.
		// So, we:
		//   - copy the existing hook details
		//   - set 'oldfunc' of the existing hook to the new 'func'
		// so that the call tree next time the SCE func is called will be:
		//   - first hook
		//     - this hook
		//       - SCE func
		// THIS ASSUMES THAT ALL HOOK FUNCS WILL CALL THROUGH TO THE ORIG FUNC

#if DEBUG_APIHOOK==1
		////sioPrint("Copy hook record from number ");
		////sioPrintHex(lhookid);
		////sioPrint("\r\n");
#endif

		int lnewhookid = numHookRecords++;
		memcpy(&(hookRecords[lnewhookid]), &(hookRecords[lhookid]), sizeof(HookRecord));
		hookRecords[lhookid].oldfunc = func;
		hookRecords[lnewhookid].syscall = 0;

		lhookid = lnewhookid;
	}

	sceKernelSignalSema(hookSema, 1);

#if DEBUG_APIHOOK==1
	////sioPrint("Return hookid ");
	////sioPrintHex(lhookid);
	////sioPrint("\r\n");
#endif

	return lhookid;
}

u32 apiHookByNid(const char *modname, const char *library, u32 nid, void *func)
{
	sceKernelWaitSema(hookSema, 1, 0);

	if (numHookRecords == MAX_HOOK_RECORDS)
	{
		sceKernelSignalSema(hookSema, 1);
		return -1;
	}

	u32 lhookid = -1;

  /***************************************************************************/
	/* First, let's see if we already have a hook for this function.           */
	/***************************************************************************/
	int ii;
	for (ii = 0; ii < numHookRecords; ii++)
	{
		if ((hookRecords[ii].nid == nid) &&
				(strcmp(hookRecords[ii].libname, library) == 0) &&
				(strcmp(hookRecords[ii].modname, modname) == 0))
		{
			lhookid = ii;
		}
	}

	if (lhookid == -1)
	{
		// no existing hook, so this is simple.
		SceModule *pMod = sceKernelFindModuleByName(modname);
		//pspSdkFixupImports(pMod->modid);
		// ^ Crashes PSP
		if (pMod)
		{
			u32 addr = libsFindExportByNid(pMod->modid, library, nid);
			if(addr)
			{
				void * syscalladdr = find_syscall_addr(addr);
				if(!apiHookAddr(syscalladdr, func))
				{
					////sioPrint("Failed to apiHookAddr\n");
					//ShowOutput("Failed to apiHookAddr");
				}
				else
				{
					// success hooking.  Record the details.
					ShowOutput("Successfully Hooked");
					lhookid = numHookRecords++;
					hookRecords[lhookid].oldfunc = (void*)addr;
					strcpy(hookRecords[lhookid].libname, library);
					strcpy(hookRecords[lhookid].modname, modname);
					hookRecords[lhookid].nid     = nid;
					hookRecords[lhookid].func    = func;
					hookRecords[lhookid].syscall = syscalladdr;
				}
			}
			else
			{
				ShowOutput("Failed to find exportbynid");
				////sioPrint("Failed to find exportbynid\n");
			}
		}
		else
		{
			ShowOutput("Failed to find modulebyname");
		}
	}
	else
	{
		// We have an existing hook for this function.
		// So, we:
		//   - copy the existing hook details
		//   - set 'oldfunc' of the existing hook to the new 'func'
		// so that the call tree next time the SCE func is called will be:
		//   - first hook
		//     - this hook
		//       - SCE func
		// THIS ASSUMES THAT ALL HOOK FUNCS WILL CALL THROUGH TO THE ORIG FUNC

#if DEBUG_APIHOOK==1
		////sioPrint("Copy hook record from number ");
		////sioPrintHex(lhookid);
		////sioPrint("\r\n");
#endif

		int lnewhookid = numHookRecords++;
		memcpy(&(hookRecords[lnewhookid]), &(hookRecords[lhookid]), sizeof(HookRecord));
		hookRecords[lhookid].oldfunc = func;
		hookRecords[lnewhookid].syscall = 0;

		lhookid = lnewhookid;
	}

	sceKernelSignalSema(hookSema, 1);

#if DEBUG_APIHOOK==1
	////sioPrint("Return hookid ");
	////sioPrintHex(lhookid);
	////sioPrint("\r\n");
#endif

	return lhookid;
}

typedef u32 (*HOOKFUNC)(u32, u32, u32, u32, u32, u32, u32);

u32 callApiHook(u32 hookid, ...)
{
	int lreturn;
	u32 args[7];
	int ii;
	va_list val;

	// save the argument list
	va_start(val, hookid);
	for (ii = 0; ii < 7; ii++)
	{
		args[ii] = va_arg(val, u32);
	}
	va_end(val);

#if DEBUG_APIHOOK==1
	////sioPrint("call api hook #");
	////sioPrintHex(hookid);
	////sioPrint("\r\n");
#endif

	sceKernelWaitSema(hookSema, 1, 0);

	// find the hook
	if ((hookid >= 0) && (hookid < numHookRecords))
	{
		HOOKFUNC lfunc = (HOOKFUNC) (hookRecords[hookid].oldfunc);

		sceKernelSignalSema(hookSema, 1);

		if (lfunc)
		{
#if DEBUG_APIHOOK==1
			//sioPrint("Calling with args: \r\n");
			for (ii = 0; ii < 2; ii++)
			{
				//sioPrintWord(args[ii]);
				//sioPrint("\r\n");
			}
#endif

			lreturn = lfunc(args[0], args[1], args[2], args[3], args[4], args[5], args[6]);

#if DEBUG_APIHOOK==1
			//sioPrint("Return: ");
			//sioPrintWord(lreturn);
			//sioPrint("\r\n");
#endif
		}
		else
		{
			lreturn = -2;
		}
	}
	else
	{
		sceKernelSignalSema(hookSema, 1);

		lreturn = -1;
	}

	return lreturn;
}

void unHookApi(u32 hookid)
{
	sceKernelWaitSema(hookSema, 1, 0);

	// find the hook
	if ((hookid >= 0) && (hookid < numHookRecords))
	{

		/* Restore original function */
		if(hookRecords[hookid].oldfunc)
		{
			u32 *lsyscall = hookRecords[hookid].syscall;

			if (lsyscall)
			{
				*lsyscall = (u32) hookRecords[hookid].oldfunc;
				hookRecords[hookid].oldfunc = NULL;
				sceKernelDcacheWritebackInvalidateRange(lsyscall, sizeof(void *));
				sceKernelIcacheInvalidateRange(lsyscall, sizeof(void *));
			}
		}
	}

	sceKernelSignalSema(hookSema, 1);
}
